import pygame
import random
import numpy as np

pygame.init()

WIDTH, HEIGHT = 400, 400
GRID_SIZE = 20

WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (213, 50, 80)
GREEN = (0, 255, 0)
BLUE = (50, 153, 213)

screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Snakes")


class SnakeGame:
    def __init__(self):
        self.game_over = None
        self.score = None
        self.food = None
        self.snake_dir = None
        self.snake = None
        self.reset()

    def reset(self):
        self.snake = [(WIDTH // 2, HEIGHT // 2)]  # 蛇的初始位置
        self.snake_dir = (GRID_SIZE, 0)  # 初始方向：向右
        self.food = None
        self.score = 0
        self.game_over = False
        self.spawn_food()  # 生成初始食物

        return self.get_state()

    def spawn_food(self):
        self.food = (random.randrange(0, WIDTH, GRID_SIZE), random.randrange(0, HEIGHT, GRID_SIZE))

    def get_state(self):
        if self.game_over:
            return np.zeros(14)  # 游戏结束时返回零状态

        snake_head = self.snake[0]
        food_distance = np.array([self.food[0] - snake_head[0], self.food[1] - snake_head[1]])

        dist_to_left_wall = snake_head[0]
        dist_to_right_wall = WIDTH - snake_head[0] - GRID_SIZE
        dist_to_top_wall = snake_head[1]
        dist_to_bottom_wall = HEIGHT - snake_head[1] - GRID_SIZE

        MAX_DIST = 400

        dist_to_left_body = MAX_DIST
        dist_to_right_body = MAX_DIST
        dist_to_top_body = MAX_DIST
        dist_to_bottom_body = MAX_DIST

        for segment in self.snake[1:]:
            dist_to_segment_x = snake_head[0] - segment[0]
            dist_to_segment_y = snake_head[1] - segment[1]

            if dist_to_segment_x > 0:
                dist_to_left_body = min(dist_to_left_body, dist_to_segment_x)
            else:
                dist_to_right_body = min(dist_to_right_body, -dist_to_segment_x)

            if dist_to_segment_y > 0:
                dist_to_top_body = min(dist_to_top_body, dist_to_segment_y)
            else:
                dist_to_bottom_body = min(dist_to_bottom_body, -dist_to_segment_y)

        direction = [
            int(self.snake_dir == (GRID_SIZE, 0)),  # 右
            int(self.snake_dir == (-GRID_SIZE, 0)),  # 左
            int(self.snake_dir == (0, GRID_SIZE)),  # 下
            int(self.snake_dir == (0, -GRID_SIZE))  # 上
        ]

        return np.concatenate([
            food_distance,
            [dist_to_left_wall, dist_to_right_wall, dist_to_top_wall, dist_to_bottom_wall],
            [dist_to_left_body, dist_to_right_body, dist_to_top_body, dist_to_bottom_body],
            direction
        ])

    def step(self, action):
        if self.game_over:
            return self.get_state(), -10, True  # 游戏结束惩罚

        if action == 0:
            self.snake_dir = (GRID_SIZE, 0)  # 向右
        elif action == 1:
            self.snake_dir = (-GRID_SIZE, 0)  # 向左
        elif action == 2:
            self.snake_dir = (0, GRID_SIZE)  # 向下
        elif action == 3:
            self.snake_dir = (0, -GRID_SIZE)  # 向上

        new_head = (self.snake[0][0] + self.snake_dir[0], self.snake[0][1] + self.snake_dir[1])

        if new_head[0] < 0 or new_head[0] >= WIDTH or new_head[1] < 0 or new_head[
            1] >= HEIGHT or new_head in self.snake:
            self.game_over = True
            return self.get_state(), -10, True

        self.snake = [new_head] + self.snake

        reward = 1.0  # 基础生存奖励

        reward -= 2
        if new_head == self.food:
            self.score += 1
            reward = 50
            self.spawn_food()
        else:
            prev_dist = np.linalg.norm(np.array(self.food) - np.array(self.snake[1]))
            new_dist = np.linalg.norm(np.array(self.food) - np.array(new_head))
            reward += (prev_dist - new_dist) * 0.5 / GRID_SIZE  # 归一化
            self.snake.pop()

        if (new_head[0] <= 0 or new_head[0] >= WIDTH or
                new_head[1] <= 0 or new_head[1] >= HEIGHT):
            reward -= 15
            self.game_over = True

        return self.get_state(), reward, False

    def render(self):
        screen.fill(BLACK)
        for segment in self.snake:
            pygame.draw.rect(screen, GREEN, pygame.Rect(segment[0], segment[1], GRID_SIZE, GRID_SIZE))
        pygame.draw.rect(screen, RED, pygame.Rect(self.food[0], self.food[1], GRID_SIZE, GRID_SIZE))
        pygame.display.update()
